package amrwriter

import (
	"common/rtpengine/rtp"
	"errors"
	"fmt"
	"io"
	"os"
)

const (
	MediaAmr = iota
	MediaAmrOctet
	MediaAmrWb
	MediaAmrWbOctet

	OctetAlignedMode       = 1
	BandwidthEfficientMode = 2
)

var AmrFrameSize = map[int][]int{MediaAmr: {12, 13, 15, 17, 19, 20, 26, 31, 5, 0, 0, 0, 0, 0, 0, 0}, MediaAmrWb: {17, 23, 32, 36, 40, 46, 50, 58, 60, 5, 5, 0, 0, 0, 0, 0}}
var AmrFrameBitSize = map[int][]int{MediaAmr: {95, 103, 118, 134, 148, 159, 204, 244, 39, 43, 38, 37, 0, 0, 0, 0}, MediaAmrWb: {132, 177, 253, 285, 371, 365, 397, 461, 477, 40, 0, 0, 0, 0, 0, 0}}

type (
	// AmrWbWriter is used to take RTP packets, parse them and
	// write the data to an io.Writer.
	AmrWbWriter struct {
		writer      io.Writer
		hasKeyFrame bool
		mode        int
		mediaCode   int

		innerBuf []byte
	}
)

// NewWb New builds a new H264 writer
func NewWb(filename string) (*AmrWbWriter, error) {
	f, err := os.Create(filename)
	if err != nil {
		return nil, err
	}

	return NewWbWith(f), nil
}

// NewWbWith NewWith initializes a new H264 writer with an io.Writer output
func NewWbWith(w io.Writer) *AmrWbWriter {
	return &AmrWbWriter{
		writer:    w,
		mode:      OctetAlignedMode,
		mediaCode: MediaAmrWb,
		innerBuf:  make([]byte, 256),
	}
}

// WriteRTP adds a new packet and writes the appropriate headers for it
func (f *AmrWbWriter) WriteRTP(packet *rtp.Packet) error {
	if len(packet.Payload) == 0 {
		return nil
	}

	n, data, err := f.ExtractSpeechFrame(packet.Payload)
	if err != nil || n == 0 {
		return err
	}

	_, err = f.writer.Write(data)

	return err
}

// WriteRTP adds a new packet and writes the appropriate headers for it
func (f *AmrWbWriter) WriteFile(payload []byte) error {
	if len(payload) == 0 {
		return nil
	}
	_, err := f.writer.Write(payload)

	return err
}

// Close closes the underlying writer
func (f *AmrWbWriter) Close() error {
	if f.writer != nil {
		if closer, ok := f.writer.(io.Closer); ok {
			return closer.Close()
		}
	}

	return nil
}

func (f *AmrWbWriter) ExtractSpeechFrame(payload []byte) (int, []byte, error) {
	if len(payload) < 2 {
		return 0, nil, errors.New(fmt.Sprintf("rtp payload is too short: %v", len(payload)))
	}
	if f.mode == OctetAlignedMode {
		return f.handleOaMode(payload)
	}
	return f.handleBeMode(payload)
}

func (f *AmrWbWriter) handleBeMode(payload []byte) (int, []byte, error) {
	// packing frame with TOC: frame type and quality bit
	// RTP=[CMR(4bit)[F][FT(4bit)][Q][..speechFrame]] -> storage=[0][FT(4bit)][Q][0][0]
	// cmr := (rtpFrameHeader[0] & 0xF0) >> 4

	rtpFrameHeader := payload[0:]

	isLastFrame := (rtpFrameHeader[0]&0x08)>>3 == 0x00
	frameType := (rtpFrameHeader[0]&0x07)<<1 | (rtpFrameHeader[1]&0x80)>>7
	quality := (rtpFrameHeader[1] & 0x40) == 0x40

	fmt.Printf("Bandwidth efficient, lastFrame: %v, frameType: %d, quality: %v\n", isLastFrame, frameType, quality)

	// 多帧处理
	if !isLastFrame {
		return f.handleMultiFrames(payload)
	}

	if !quality {
		return 0, nil, errors.New("rtp payload contains a bad frame")
	}
	if AmrFrameSize[f.mediaCode][frameType] == 0 {
		return 0, nil, errors.New("amr frameSize is invalid")
	}
	// 单帧处理
	end := 1
	speechFrameHeader := (rtpFrameHeader[0]&0x07)<<4 | (rtpFrameHeader[1]&0x80)>>4
	speechFrameHeader = speechFrameHeader | (rtpFrameHeader[1]&0x40)>>4
	speechFrameSize := AmrFrameSize[f.mediaCode][frameType]
	f.innerBuf[0] = speechFrameHeader
	speechPayload := rtpFrameHeader[1:]
	if len(speechPayload) < speechFrameSize {
		return 0, nil, errors.New(fmt.Sprintf("speechFrame size [%v] does not match frameType [%v], expect: %v", len(speechPayload), frameType, speechFrameSize))
	}
	speechFrame := f.innerBuf[end:]
	for k := 0; k < speechFrameSize; k++ {
		speechFrame[k] = (speechPayload[k] & 0x3F) << 2
		if k+1 < speechFrameSize && len(speechPayload) > k+1 {
			speechFrame[k] = speechFrame[k] | (speechPayload[k+1]&0xC0)>>6
		}
	}
	end += speechFrameSize
	return 1, f.innerBuf[:end], nil
}

func (f *AmrWbWriter) handleOaMode(payload []byte) (int, []byte, error) {
	// payload header := [CMR(4bit)[R(4bit)][ILL(4bit)(opt)][ILP(4bit)(opt)]，here ignore ILL and interleaving、crc in session
	// TOC := [F][FT(4bit)][Q][P][P]
	// cmr := (rtpFrameHeader[0] & 0xF0) >> 4

	rtpFrameHeader := payload[0:]

	isLastFrame := (rtpFrameHeader[1]&0x80)&0x80 == 0x00
	frameType := (rtpFrameHeader[1] & 0x78) >> 3
	quality := (rtpFrameHeader[1]&0x04)&0x04 == 0x04

	fmt.Printf("Octet-aligned, lastFrame: %v, frameType: %d, quality: %v\n", isLastFrame, frameType, quality)

	// 多帧处理
	if !isLastFrame {
		return f.handleMultiFrames(payload)
	}

	if !quality {
		return 0, nil, errors.New("rtp payload contains a bad frame")
	}
	if AmrFrameSize[f.mediaCode][frameType] == 0 {
		return 0, nil, errors.New("amr frameSize is invalid")
	}
	// 单帧处理
	end := 1
	speechFrameHeader := frameType << 3
	speechFrameHeader = speechFrameHeader | (rtpFrameHeader[1] & 0x04)
	speechFrameSize := AmrFrameSize[f.mediaCode][frameType]
	f.innerBuf[0] = speechFrameHeader
	if len(rtpFrameHeader[2:]) < speechFrameSize {
		return 0, nil, errors.New(fmt.Sprintf("speechFrame size [%v] does not match frameType [%v], expect: %v", len(rtpFrameHeader[2:]), frameType, speechFrameSize))
	}
	copy(f.innerBuf[end:], rtpFrameHeader[2:2+speechFrameSize])
	end += speechFrameSize
	return 1, f.innerBuf[:end], nil
}

func (f *AmrWbWriter) handleMultiFrames(payload []byte) (int, []byte, error) {
	offset := 4
	octetAligned := false
	if f.mode == OctetAlignedMode {
		offset = 8
		octetAligned = true
	}

	// TODO: 多帧情况下内存复用
	var frames []byte
	var tocItems []byte
	for {
		frame, idx := Resolve(payload, offset, Round(6, octetAligned))
		if idx <= 0 {
			break
		}
		if len(frame) == 0 {
			return 0, nil, errors.New("rtp payload contains a bad multiFrame")
		}
		tocItems = append(tocItems, frame[0]&0x7c)
		offset = idx
		if frame[0]&0x80 == 0 {
			break
		}
	}
	nFrames := 0
	for _, toc := range tocItems {
		quality := (toc & 0x04) >> 2
		if quality != 1 {
			// 坏帧，忽略
			break
		}
		frameType := toc >> 3
		frames = append(frames, toc)
		if AmrFrameBitSize[f.mediaCode][frameType] == 0 {
			return 0, nil, errors.New("multi amr frameSize is invalid")
		}
		frame, idx := Resolve(payload, offset, Round(AmrFrameBitSize[f.mediaCode][frameType], octetAligned))
		offset = idx
		frames = append(frames, frame...)
		nFrames++
	}
	fmt.Printf("Handle %d multiFrame", nFrames)
	return nFrames, frames, nil
}

func Resolve(data []byte, offset, n int) ([]byte, int) {
	res := make([]byte, 0)
	if n < 0 || n+offset > 8*len(data) {
		return res, offset
	}
	if n == 0 {
		return res, offset
	}
	i := offset >> 3
	sh := uint(offset % 8)
	x := data[i]
	for j := 0; j < (n+7)>>3; j++ {
		value := (x << sh) & 0xFF
		i++
		if i < len(data) {
			x = data[i]
			value |= x >> (8 - sh)
		}
		res = append(res, value)
	}
	offset += n
	nb := uint(8 - n%8)
	if nb < 8 {
		mask := 0x100 - (1 << nb)
		res[len(res)-1] = byte(int(res[len(res)-1]) & mask)
	}
	return res, offset
}

func Round(n int, mode bool) int {
	if mode {
		n = (n + 7) >> 3 << 3
	}
	return n
}
